import re
from core.utils.logging_handler import LoggingHandler

class QueryExecutor:
    def __init__(self, query_code):
        self.logger = LoggingHandler().get_logger()
        self.logger.info("初始化 QueryExecutor")
        self.query_code = self._extract_and_sanitize_code(query_code)

    def _extract_and_sanitize_code(self, code):
        # 提取出在 ```python ``` 和 ``` 之间的内容
        code_match = re.search(r"```python(.*?)```", code, re.DOTALL)
        if code_match:
            code = code_match.group(1).strip()
        else:
            # 如果没有找到代码块，假设整个代码字符串都是代码
            code = code.strip()

        # 保留缩进，移除行内和块注释
        sanitized_lines = []
        for line in code.split("\n"):
            # 保留行首的空白（缩进）
            line_without_comment = re.sub(r"#.*", "", line).rstrip()
            if line_without_comment:
                sanitized_lines.append(line_without_comment)

        # 将代码行重新组合
        sanitized_code = "\n".join(sanitized_lines)
        return sanitized_code

    def execute(self, df):
        self.logger.info("开始执行查询代码")
        local_env = {'df': df, 'result': None}
        self.logger.info(self.query_code)
        try:
            exec(self.query_code, {}, local_env)
            result = local_env.get('result')
            self.logger.info("查询代码执行成功")
        except Exception as e:
            self.logger.error(f"执行查询代码时发生异常: {e}", exc_info=True)
            result = f"Execution error: {e}"
        
        return result



